# -*- coding: utf-8 -*-
"""
Created on Wed Dec 16 09:50:27 2020

@author: Team317
"""
import cv2
import os
import random
import torch
import numpy as np
import shutil
from PIL import Image

# 处理从github上得到的数据集
def dealImage(root_path='../AFAD-Full', save_root='../AFAD_image'):
    if not os.path.exists(save_root):
        os.mkdir(save_root)
    age_set = os.listdir(root_path)
    for first_path in age_set:  # 遍历每一个年龄
        # count = 0   # 给每个图像集统计图片数量
        age_folder = str(int((int(first_path) - 5) / 10) + 2)  # 计算所在年龄段
        age_path = os.path.join(root_path, first_path)
        
        #创建每个年龄的存储文件夹
        save_age_path = os.path.join(save_root, age_folder)
        if not os.path.exists(save_age_path):
            os.mkdir(save_age_path)
             
        gender_set = os.listdir(age_path)
        for second_path in gender_set:  # 遍历两个性别
            gender_path = os.path.join(age_path, second_path)
            image_set = os.listdir(gender_path)
            
   
            for image_file in image_set:  # 遍历每一张图片
                if image_file != 'Thumbs.db':
                    # 读取图片并缩放为112，112
                    image_path = os.path.join(gender_path, image_file)
                    image = cv2.imread(image_path)
                    try:
                        deal_image = cv2.resize(image,(112,112))
                    except:
                        print(image_path)
                        
                    # 保存图片：方式一
                    # count+=1
                    # image_name = '{0:05d}.jpg'.format(count)
                    # save_image_path = os.path.join(save_age_path, image_name)
                    # cv2.imwrite(save_image_path, deal_image)
                    
                    # 保存图片：方式二
                    postfix = len(os.listdir(save_age_path))  # 用图片集的大小作为后缀
                    image_name = '{0:05d}.jpg'.format(postfix)
                    save_image_path = os.path.join(save_root, age_folder, image_name)
                    cv2.imwrite(save_image_path, deal_image)
                        
            
            # print("年龄为{}的图片数量:{}".format(age_folder,count))

# 处理原始数据集
def dealImage2(root_path='../age_image/img_face2', save_root='../AFAD_image1'):
    
    age_set = os.listdir(root_path)
    age_set1 = age_set[5:]
    for age_folder in age_set1:
        age_path = os.path.join(root_path, age_folder)
        image_set = os.listdir(age_path)
        
        # 创建保存图片的文件夹
        save_age_path = os.path.join(save_root, age_folder[-1])
        if not os.path.exists(save_age_path):
            os.mkdir(save_age_path)
        postfix = len(os.listdir(save_age_path))  # 图片保存时的后缀数字
        
        for image_name in image_set:
            image_path = os.path.join(age_path, image_name)
            
            
            # 将尾数为偶数的图片保存下来
            last_num = int(image_name[-9]) 
            if last_num%2 == 0:
                image = cv2.imread(image_path)
                deal_image = cv2.resize(image,(112,112))
                
                image_save_name = '{0:05d}.jpg'.format(postfix)
                image_save_path = os.path.join(save_age_path, image_save_name)
                cv2.imwrite(image_save_path, deal_image)
                
                postfix += 1  # 将后缀加一
                
        print("第{}个年龄段已完成".format(age_folder[-1]))
        
        
# 统计每个年龄段的图片数量
def fileCount(root_path = '../AFAD_image1'):
    age_set = os.listdir(root_path)
    for age_folder in age_set:
        age_path = os.path.join(root_path, age_folder)
        image_set = os.listdir(age_path)

        print("文件夹{}中的图片数量为：{}".format(age_folder, len(image_set)))
# 生成图片总集
def generateTotalSet(root_path='../AFAD_image1', save_path='../totalSet'):
    # 生成总的保存目录
    if not os.path.exists(save_path):
        os.mkdir(save_path)
        
    age_set = os.listdir(root_path)
    for age_folder in age_set:
        age_path = os.path.join(root_path, age_folder)
        image_set = os.listdir(age_path)
        count = 0
        save_age_path = os.path.join(save_path, age_folder)
        if not os.path.exists(save_age_path):
            os.mkdir(save_age_path)
            
        for image_name in image_set:
            
            # 最多保存5000张训练图片
            if count < 5000:
                # 图片读取路径
                image_path = os.path.join(age_path, image_name) 
                image = cv2.imread(image_path)
                # 图片保存路径
                image_save_path = os.path.join(save_path, age_folder,image_name)
                cv2.imwrite(image_save_path, image)
            else:
                break
            
            count += 1
            


# 生成训练和测试集
def generateSet(root_path='../trainSet', train_save_path='../trainSet', test_save_path='../testSet'):
    if not os.path.exists(train_save_path):
        os.mkdir(train_save_path)
    if not os.path.exists(test_save_path):
        os.mkdir(test_save_path)
        
    age_set = os.listdir(root_path)
    for age_folder in age_set:
        age_path = os.path.join(root_path, age_folder)
        image_set = os.listdir(age_path)
        test_age_save_path = os.path.join(test_save_path, age_folder)
        if not os.path.exists(test_age_save_path):
            os.mkdir(test_age_save_path)
        # 移动图片
        count = 0
        for image_name in image_set:
            # 移动100张图片
            if True:
                
                src_path = os.path.join(root_path, age_folder, image_name)
                image = cv2.imread(src_path)
                image_name = '{0:05d}.jpg'.format(count)
                dst_path = os.path.join(test_save_path, age_folder, image_name)
                cv2.imwrite(dst_path, image)
                # shutil.move(src_path, dst_path)
            else:
                break
            
            count += 1
            
    
        
    
# 尝试生成网络输入图片       
def getImage(root_path = '../AFAD_image', batch=8):
    label_list = [random.randint(0,9) for i in range(0,batch)]
    transform = transforms.Compose([
            transforms.Resize(112),
            transforms.ColorJitter(0.2,0.2,0.2,0.01),
            transforms.ToTensor(),
            # transforms.ColorJitter()
            
            # transforms.Normalize((0.5, 0.5, 0.5), (128 / 255., 128 / 255., 128 / 255.))
        ])
    image_list = []
    path_list = [] 
    for label in label_list:
        age_set_path = os.path.join(root_path, str(label))
        set_len = len(os.listdir(age_set_path))
        image_index = random.randint(1,set_len-1)
        
        image_path = os.path.join(age_set_path, '{0:05d}.jpg'.format(image_index))
        image = Image.open(image_path)
        image = transform(image)
        # image = image.transpose(1, -1)
        # image_list = image_list.append(image)
        # image = cv2.imread(image_path)
        # image = cv2.resize(image, (112,112))
        
        image_list.append(image)
        path_list.append(image_path)
        
    # image_set1 = np.array(image_list,dtype='float16')
    # image_set1 = np.array(image_list)
    # image_set2 = torch.tensor(image_list)
    image_set2 = torch.stack(image_list)
    image_set2.type = torch.cuda.FloatTensor
    # image_set3 = image_set2.permute(0,-1,2,1)
    return image_set2, torch.tensor(label_list), path_list
    

from torchvision import transforms as transforms
import matplotlib.pyplot as plt
if "__main__" == __name__:
    root_path = os.path.join('..','AFAD-Full')
    save_root = os.path.join('..','AFAD_image1')
    # dealImage(root_path, save_root)
    # fileCount()
    # dealImage2()
    # generateTotalSet()
    # generateSet()
    # fileCount('../testSet')
    
    # info_path = 'image_info.txt'
    # imageCount(info_path)
    # result = getImage(save_root)
    # print(result)
    
    root_path = os.path.join('..','trainSet1')
    # train_path = os.path.join("..", "trainSet1")
    # generateSet(root_path,test_save_path = train_path)
    
    
    # 展示图片
    batch_size = 8
    image_set, label_set, image_path= getImage(root_path, batch=batch_size)
    width = batch_size / 4
    fig = plt.figure("image")
    for i in range(len(image_set)):
        plt.subplot(4, width, i+1)
        
        image = np.array(image_set[i])
        plt.imshow(image)
    
    plt.show()
    
    
    # cv2.imshow("hh", image)
    # transforms.ColorJitter(brightness=1)(image_set[0])
    